# -*- coding: utf-8 -*-
from flask import Flask, request, jsonify
import numpy as np
import json
from scipy.signal import butter, filtfilt
from scipy.signal import hilbert
from PyEMD import EEMD, EMD, Visualisation
from vmdpy import VMD
from scipy.fftpack import fft
import scipy.signal
from scipy.signal import stft

app = Flask(__name__)


def butter_lowpass(cutoff, fs, order=5):
    nyq = 0.5 * fs  # 奈奎斯特频率
    normal_cutoff = cutoff / nyq  # 归一化截止频率
    b, a = butter(order, normal_cutoff, btype='low', analog=False)
    return b, a


def lowpass_filter(data, cutoff, fs, order=5):
    data = np.ravel(data)  # 确保数据是一维的
    b, a = butter_lowpass(cutoff, fs, order=order)
    y = filtfilt(b, a, data)
    return y


def demean(signal):
    # 计算 filtered_data 的均值
    mean_value = np.mean(signal)

    # 将每个值减去均值
    data_centered = signal - mean_value
    return data_centered


# 信号分解
def decompose_lw(signal, t, method='vmd', K=3, draw=1):
    names = ['emd', 'eemd', 'vmd']
    idx = names.index(method)

    # EMD 分解
    if idx == 0:
        emd = EMD()
        IMFs = emd.emd(signal)

    # VMD 分解
    elif idx == 2:
        alpha = 2000  # moderate bandwidth constraint
        tau = 0  # noise-tolerance
        DC = 0  # no DC part imposed
        init = 1  # initialize omegas uniformly
        tol = 1e-7
        IMFs, _, _ = VMD(signal, alpha, tau, K, DC, init, tol)

    # EEMD 分解
    else:
        eemd = EEMD()
        emd = eemd.EMD
        emd.extrema_detection = "parabol"
        IMFs = eemd.eemd(signal, t)
    return IMFs


# 希尔伯特变换及画时频谱
def hhtlw(IMFs, t, f_range=[0, 100], t_range=[0, 20], ft_size=[128, 128], draw=1):
    fmin, fmax = f_range[0], f_range[1]  # 时频图所展示的频率范围
    tmin, tmax = t_range[0], t_range[1]  # 时间范围
    fdim, tdim = ft_size[0], ft_size[1]  # 时频图的尺寸（分辨率）
    dt = (tmax - tmin) / (tdim - 1)
    df = (fmax - fmin) / (fdim - 1)

    vis = Visualisation()

    # 希尔伯特变换
    c_matrix = np.zeros((fdim, tdim))
    for imf in IMFs:
        imf = np.array([imf])
        # 瞬时频率
        freqs = abs(vis._calc_inst_freq(imf, t, order=False, alpha=None))
        # 瞬时幅值
        amp = abs(hilbert(imf))
        # 去掉为1的维度
        freqs = np.squeeze(freqs)
        amp = np.squeeze(amp)
        # 转换成矩阵
        temp_matrix = np.zeros((fdim, tdim))
        n_matrix = np.zeros((fdim, tdim))
        for i, j, k in zip(t, freqs, amp):
            if tmin <= i <= tmax and fmin <= j <= fmax:
                temp_matrix[round((j - fmin) / df)][round((i - tmin) / dt)] += k
                n_matrix[round((j - fmin) / df)][round((i - tmin) / dt)] += 1
        n_matrix = n_matrix.reshape(-1)
        idx = np.where(n_matrix == 0)[0]
        n_matrix[idx] = 1
        n_matrix = n_matrix.reshape(fdim, tdim)
        temp_matrix = temp_matrix / n_matrix
        c_matrix += temp_matrix

    t = np.linspace(tmin, tmax, tdim)
    f = np.linspace(fmin, fmax, fdim)
    return t, f, c_matrix


# 快速傅里叶变换
def fft_transform(IMFs, fs):
    fft_results = []
    freqs = np.fft.fftfreq(len(IMFs[0]), 1 / fs)

    for imf in IMFs:
        fft_result = np.abs(fft(imf))
        fft_results.append(fft_result)

    return freqs, fft_results


# 快速傅里叶变换
@app.route('/fft', methods=['POST'])
def myFft():
    # 获取 JSON 数据
    data = request.get_json()
    # 获取 startList 数组
    data_list = data.get('list', [])
    if data_list is None:
        return
    # 确保 data_list 是一维数组
    data_array = np.array(data_list)

    fs = 1000  # 采样频率为 1000 Hz
    cutoff = 200  # 截止频率为 200 Hz
    # 低通滤波
    filtered_data = lowpass_filter(data_array, cutoff, fs)

    # 去均值
    data_centered = demean(filtered_data)
    # 数据点数量
    N = len(data_centered)
    # 进行快速傅里叶变换
    fft_result = np.fft.fft(data_centered)
    freqs = np.fft.fftfreq(N, d=1 / fs)

    positive_freqs = freqs[:N // 2]  # 只保留正频率部分
    amplitude_spectrum = np.abs(fft_result[:N // 2]) / N  # 归一化幅度谱

    # 仅保留0到100 Hz范围内的频率和幅度谱
    freq_range_mask = (positive_freqs <= 200)
    filtered_freqs = positive_freqs[freq_range_mask]
    filtered_amplitude_spectrum = amplitude_spectrum[freq_range_mask]

    # 时间轴转换为毫秒
    time_axis = np.arange(N) / fs  # 计算时间轴（秒）
    data = [
        {
            "name": "压力脉动",
            "x": time_axis.tolist(),
            "y": [data_centered.tolist()]
        },
        {
            "name": "快速傅里叶变换",
            "x": filtered_freqs.tolist(),
            "y": [filtered_amplitude_spectrum.tolist()]
        }
    ]
    return json.dumps(data)


# EMD\EEMD\VMD+Hilbert时频图
@app.route('/salvageHilbert', methods=['POST'])
def salvageHilbert():
    # 获取 JSON 数据
    data = request.get_json()
    # 21VMD\22EMD\23EEMD,获取类型
    type = data.get('type', int)
    method = ['vmd','emd','eemd']

    # 获取 startList 数组
    data_list = data.get('list', [])
    # 确保 data_list 是一维数组
    signal = np.array(data_list)
    fs = 1000  # 采样频率
    t = np.arange(0, len(signal) / fs, 1 / fs)
    # 低通滤波
    filtered_data = lowpass_filter(signal, 100, fs)

    # 去均值
    data_centered = demean(filtered_data)
    # 信号分解
    IMFs = decompose_lw(data_centered, t, method=method[type-21], K=5)
    # 生成时频图
    tt, ff, c_matrix = hhtlw(IMFs, t, f_range=[0, 100], t_range=[0, 20], ft_size=[128, 128])

    data_list = []
    y_list = []
    for i in range(len(IMFs)):
        y_list.append(IMFs[i].tolist())
    data_item = {
        "name": "分解信号",
        "x": t.tolist(),
        "y": y_list
    }
    data_list.append(data_item)

    data_item = {}
    data_item["name"] = "希尔伯特谱"
    data_item["x"] = tt.tolist()
    data_item["y"] = [ff.tolist()]
    data_item["matrix"] = c_matrix.tolist()
    data_list.append(data_item)
    return json.dumps(data_list)


# EMD\EEMD\VMD+快速傅里叶变换
@app.route('/salvageFft', methods=['POST'])
def salvageFft():
    # 获取 JSON 数据
    data = request.get_json()
    # 31VMD\32EMD\33EEMD,获取类型
    type = data.get('type', int)
    method = ['vmd', 'emd', 'eemd']
    # 获取 startList 数组
    data_list = data.get('list', [])
    # 确保 data_list 是一维数组
    signal = np.array(data_list)
    fs = 1000  # 采样频率
    t = np.arange(0, len(signal) / fs, 1 / fs)
    # 低通滤波
    filtered_data = lowpass_filter(signal, 100, fs)
    # 去均值
    data_centered = demean(filtered_data)
    # 信号分解
    IMFs = decompose_lw(data_centered, t, method=method[type-31], K=4)
    # 对IMFs应用快速傅里叶变换
    freqs, fft_results = fft_transform(IMFs, fs)

    data_list = []
    y_list = []
    for i in range(len(IMFs)):
        y_list.append(IMFs[i].tolist())
    data_item = {
        "name": "分解信号",
        "x": t.tolist(),
        "y": y_list
    }
    data_list.append(data_item)

    positive_freqs = freqs[:len(freqs) // 2]
    # 仅保留0到100 Hz范围内的频率和幅度谱
    freq_range_mask = (positive_freqs <= 100)
    filtered_freqs = positive_freqs[freq_range_mask]

    data_item = {}
    data_item["name"] = "快速傅里叶变换"
    data_item["x"] = filtered_freqs.tolist()
    y_list = []
    for i, fft_result in enumerate(fft_results):
        positive_amplitude_spectrum = fft_result[:len(fft_result) // 2]
        filtered_amplitude_spectrum = positive_amplitude_spectrum[freq_range_mask]
        y_list.append(filtered_amplitude_spectrum.tolist())
    data_item["y"] = y_list
    data_list.append(data_item)
    return json.dumps(data_list)


# 希尔伯特变化提取包络谱图
@app.route('/plotHilbert', methods=['POST'])
def plotHilbertTransform():
    # 获取 JSON 数据
    data = request.get_json()
    # 获取 startList 数组
    data_list = data.get('list', [])
    # 确保 data_list 是一维数组
    signal = np.array(data_list)
    fs = 1000  # 采样频率
    # 低通滤波
    filtered_signal = lowpass_filter(signal, 100, fs)
    # 去均值
    data_centered = demean(filtered_signal)
    # 进行希尔伯特变换
    analytic_signal = hilbert(data_centered)
    at = np.sqrt(analytic_signal ** 2 + data_centered ** 2)

    sampling_rate = 1000
    at = at - np.mean(at)  # 去直流分量
    am = np.fft.fft(at)  # 对希尔伯特变换后的at做fft变换获得幅值
    am = np.abs(am)  # 对幅值求绝对值（此时的绝对值很大）
    am = am / len(am) * 2
    am = am[0: int(len(am) / 2)]
    freq = np.fft.fftfreq(len(at), d=1 / sampling_rate)  # 获取fft频率，此时包括正频率和负频率
    freq = freq[0:int(len(freq) / 2)]  # 获取正频率

    # 仅保留0到100 Hz范围内的频率和幅度谱
    freq_range_mask = (freq <= 200)
    filtered_freqs = freq[freq_range_mask]
    am = am[freq_range_mask]

    data_list = []
    data_item = {
        "name": "希尔伯特变化提取包络谱图",
        "x": filtered_freqs.tolist(),
        "y": [am.tolist()]
    }
    data_list.append(data_item)
    return json.dumps(data_list)


# 功率谱
@app.route('/powerSpectrum', methods=['POST'])
def powerSpectrum():
    data = request.get_json()
    data_list = data.get('list', [])
    # 确保 data_list 是一维数组
    signal = np.array(data_list)
    fs = 1000  # 采样频率
    # 低通滤波
    filtered_signal = lowpass_filter(signal, 100, fs)
    # 去均值
    data_centered = demean(filtered_signal)
    # 计算功率谱密度（PSD）
    frequencies, psd = scipy.signal.periodogram(data_centered, fs=1000)  # fs 是采样频率，假设为1000 Hz
    # 仅保留0到100 Hz范围内的频率和幅度谱
    freq_range_mask = (frequencies <= 200)
    filtered_freqs = frequencies[freq_range_mask]
    filtered_amplitude_spectrum = psd[freq_range_mask]
    data_list = []
    data_item = {
        "name": "功率谱图",
        "x": filtered_freqs.tolist(),
        "y": [filtered_amplitude_spectrum.tolist()]
    }
    data_list.append(data_item)
    return json.dumps(data_list)


# 短时傅里叶变换
@app.route('/stft', methods=['POST'])
def myStft():
    data = request.get_json()
    data_list = data.get('list', [])
    # 确保 data_list 是一维数组
    signal = np.array(data_list)
    fs = 1000  # 采样频率
    # 低通滤波
    filtered_signal = lowpass_filter(signal, 100, fs)
    # 去均值
    data_centered = demean(filtered_signal)
    window = 'hann'
    # frame长度
    n = 64

    # STFT
    f, t, Z = stft(data_centered, fs=fs, window=window, nperseg=n)
    # 选择 0 到 300 Hz 范围的频率
    freq_mask = (f <= 300)
    f = f[freq_mask]  # 频率范围
    Z = Z[freq_mask]  # 仅保留对应的频率范围
    # 求幅值
    Z = np.abs(Z)

    data_list = []
    data_item = {}
    data_item["name"] = "短时傅里叶变换"
    data_item["x"] = t.tolist()
    data_item["y"] = [f.tolist()]
    data_item["matrix"] = Z.tolist()
    data_list.append(data_item)
    return json.dumps(data_list)

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000)
